[ARM] MLAS: portable machine-code SVE elementwise kernels, FEXPA exp - #31145
Draft
mirounga wants to merge 2 commits into
Draft
[ARM] MLAS: portable machine-code SVE elementwise kernels, FEXPA exp#31145mirounga wants to merge 2 commits into
mirounga wants to merge 2 commits into
Conversation
Rework the SVE elementwise kernels (f32: Erf, Logistic, Exp, SumExp, Softmax, LogSoftmax, ReduceMaximum, ReduceMinimumMaximum; fp16: Erf, Gelu, Tanh) so they no longer require an SVE-capable compiler or a Linux-only build, and speed them up while doing so. Structure: each kernel becomes a relocation-free *Impl function that receives its constants through table pointers (no literal pools, no global data references), fronted by a plain C++ dispatch layer (sve/elementwise_sve_dispatch.cpp) that owns the public entry points and fastpaths. The production kernels ship as portable machine code in the style of Arm's KleidiAI library (aarch64/elementwise_sve_asm.S: raw instruction words via KAI_ASM_INST, which is GAS ".inst" on Linux/macOS and armasm64 "DCD" on Windows via aarch64/kai_asm_macros.h, macro set adopted verbatim from KleidiAI). The SVE intrinsics translation units remain in-tree as the reference implementation and regeneration source, selectable with cmake option onnxruntime_SVE_ELEMENTWISE_ASM=OFF. The exp path uses the SVE FEXPA instruction (base SVE, no extension required) for the SumExp/exp evaluation, with the input clamped at -88.0f: below -127*ln2 the FEXPA index underflows into the exponent field and would produce NaN. Two kernels gain platform-dispatch coverage on AArch64/SVE that previously existed only on AMD64/RISCV64: ComputeExpF32Kernel and ReduceMinimumMaximumF32Kernel (MLAS_PLATFORM fields moved to a shared guard; call sites in compute.cpp/quantize.cpp extended with MLAS_USE_SVE). Windows enablement: runtime detection via PF_ARM_SVE_INSTRUCTIONS_AVAILABLE (SDK-#ifdef-guarded, Windows 11 24H2+), sources wired through the existing cl /P + armasm64 pipeline, the fp16 kernel dispatch un-gated from !_WIN32 (the NEON fp16 fallbacks remain POSIX-only; on Windows without SVE the routines stay null and callers use their scalar fallbacks). Structurally complete; awaiting validation on Windows-on-ARM SVE hardware. Measured on Cortex-X925/A725 (SVE VL=128), big-core pinned, 104 median cells, versus the SVE intrinsics kernels currently shipping (real-time medians, lower is better): ComputeSoftmaxInplace (32 cells) 0.772x (1.30x speedup) GeluErf fused/unfused (36 cells) 0.795-0.797x Silu (36 cells) 0.927-0.930x All 104 cells 0.832x (1.20x speedup) Correctness: full onnxruntime_mlas_test with the machine-code kernels linked: 33213 passed, 0 failed. The code is fully position-independent and vector-length agnostic (runtime cntw/cnth + whilelo loops).
mirounga
marked this pull request as draft
July 29, 2026 21:53
Follow-up to the SVE elementwise kernels, addressing the findings of Windows-on-ARM validation (GB10 / Cortex-X925+A725, MSVC + armasm64, SVE active at runtime: 0 test regressions, 1.20-1.60x on the exp/softmax/erf family on both CPU clusters, but one real regression). MlasSveReduceMinimumMaximumF32KernelImpl was measured 2.0-3.9x SLOWER than the generic NEON kernel (monotonic in N, plateauing — worst on the wide out-of-order X925), and this path feeds dynamic-quantization scale selection via MlasFindMinMaxElement. Root cause: it was the one kernel written as a naive loop — a single min/max accumulator pair processing one vector per iteration, generating a whilelt predicate and using merging (_m) operations in the hot loop — while the generic NEON kernel runs four independent accumulator pairs. (This dispatch slot did not exist on ARM before the parent commit, so the regression was introduced there, and the Linux benchmark families never covered it.) Rewritten with four unpredicated accumulator pairs and a predicated tail, matching the adjacent ReduceMaximum kernel's structure: now at parity with the generic kernel at VL=128 (1.00x for N >= 1024, was 3.9x slower at the plateau; worst case 1.12x at N = 128, sub-nanosecond), with headroom on wider-VL hardware. Numerics are bit-identical to generic, as before. Add an SVE f32 Tanh kernel: the same clamped P13/Q6 rational polynomial as the generic kernel (coefficients shared by symbol from MlasTanhConstants), in the same table-pointer Impl form as Erf/Logistic. The TanhKernelRoutine dispatch slot moves to the shared SVE/AMD64/RISCV64 platform block. Measured 0.748-0.751x of the generic kernel's time across N = 128..256K (1.33x speedup), with identical accuracy (2.15e-07 max relative error, same polynomial). Note for benchmark harnesses: BM_Tanh is no longer a negative control. EltwiseAdd was also examined and deliberately left on the generic path: it is a memory-bound streaming add, and SVE at VL=128 issues exactly the same number of 128-bit memory operations as NEON — no headroom (the validation report's 1.03x reading on it was binary-layout noise, consistent with its byte-identical fingerprint). The frozen machine code is regenerated by sve/gen_sve_asm.py, now checked in (extended for multiple source TUs, a configurable optimization level and include paths, and a section-driven parser that tolerates gcc's mid-function stack-probe labels). Regeneration fidelity: all 12 unchanged functions are byte-identical to the previously shipped assembly; the file is now 14 functions / 619 instruction words. Known numerics note carried from validation: the SVE exp path flushes to zero for inputs below ~-86.7 (the FEXPA path clamps at -88.0 to avoid exponent-field underflow), where the generic path still returns denormals. The absolute error is bounded by ~2.4e-38 and is invisible to softmax sums and to the unit tests' abs-OR-rel tolerance; documented here explicitly. Correctness: full onnxruntime_mlas_test with the new kernels linked: 33213 passed, 0 failed; the untouched benchmark families re-measured at 0.99-1.01x (byte-identical regeneration moved nothing else).
| @@ -0,0 +1,199 @@ | |||
| #!/usr/bin/env python3 | |||
| @@ -0,0 +1,199 @@ | |||
| #!/usr/bin/env python3 | |||
|
|
||
|
|
||
| def run(cmd): | ||
| result = subprocess.run(cmd, capture_output=True, text=True) |
| args = ap.parse_args() | ||
|
|
||
| symbols = args.symbols.split(",") | ||
| src_dir = os.path.dirname(os.path.abspath(args.src[0])) |
| wanted = [s for s in symbols if s not in functions] | ||
| check_relocations(obj, wanted) | ||
| for name, body in extract_functions(obj, wanted, allow_missing=True).items(): | ||
| functions[name] = body |
| args = ap.parse_args() | ||
|
|
||
| symbols = args.symbols.split(",") | ||
| src_dir = os.path.dirname(os.path.abspath(args.src[0])) |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Rework the SVE elementwise kernels (f32: Erf, Logistic, Exp, SumExp, Softmax, LogSoftmax, ReduceMaximum, ReduceMinimumMaximum; fp16: Erf, Gelu, Tanh) so they no longer require an SVE-capable compiler or a Linux-only build, and speed them up while doing so.
Structure: each kernel becomes a relocation-free *Impl function that receives its constants through table pointers (no literal pools, no global data references), fronted by a plain C++ dispatch layer (sve/elementwise_sve_dispatch.cpp) that owns the public entry points and fastpaths. The production kernels ship as portable machine code in the style of Arm's KleidiAI library (aarch64/elementwise_sve_asm.S: raw instruction words via KAI_ASM_INST, which is GAS ".inst" on Linux/macOS and armasm64 "DCD" on Windows via aarch64/kai_asm_macros.h, macro set adopted verbatim from KleidiAI). The SVE intrinsics translation units remain in-tree as the reference implementation and regeneration source, selectable with cmake option onnxruntime_SVE_ELEMENTWISE_ASM=OFF.
The exp path uses the SVE FEXPA instruction (base SVE, no extension required) for the SumExp/exp evaluation, with the input clamped at -88.0f: below -127*ln2 the FEXPA index underflows into the exponent field and would produce NaN.
Two kernels gain platform-dispatch coverage on AArch64/SVE that previously existed only on AMD64/RISCV64: ComputeExpF32Kernel and ReduceMinimumMaximumF32Kernel (MLAS_PLATFORM fields moved to a shared guard; call sites in compute.cpp/quantize.cpp extended with MLAS_USE_SVE).
Windows enablement: runtime detection via
PF_ARM_SVE_INSTRUCTIONS_AVAILABLE (SDK-#ifdef-guarded, Windows 11 24H2+), sources wired through the existing cl /P + armasm64 pipeline, the fp16 kernel dispatch un-gated from !_WIN32 (the NEON fp16 fallbacks remain POSIX-only; on Windows without SVE the routines stay null and callers use their scalar fallbacks). Structurally complete; awaiting validation on Windows-on-ARM SVE hardware.
Measured on Cortex-X925/A725 (SVE VL=128), big-core pinned, 104 median cells, versus the SVE intrinsics kernels currently shipping (real-time medians, lower is better):
ComputeSoftmaxInplace (32 cells) 0.772x (1.30x speedup)
GeluErf fused/unfused (36 cells) 0.795-0.797x
Silu (36 cells) 0.927-0.930x
All 104 cells 0.832x (1.20x speedup)
Correctness: full onnxruntime_mlas_test with the machine-code kernels linked: 33213 passed, 0 failed. The code is fully position-independent and vector-length agnostic (runtime cntw/cnth + whilelo loops).
Description
Motivation and Context